import argparse
import numpy as np
import torch
from collections import defaultdict
from tqdm import tqdm
from pathlib import Path
from zipfile import ZipFile
from transforms3d.quaternions import mat2quat

from model import PL_RelPose, keypoint_dict
from utils.reproject import reprojection_error, Pose, save_submission
from utils.metrics import reproj, add, adi, compute_continuous_auc, relative_pose_error, rotation_angular_error
from datasets import dataset_dict
from configs.default import get_cfg_defaults


@torch.no_grad()
def main(args):
    config = get_cfg_defaults()
    config.merge_from_file(args.config)

    task = config.DATASET.TASK
    dataset = config.DATASET.DATA_SOURCE
    device = args.device

    test_num_keypoints = test_num_keypoints = config.MODEL.TEST_NUM_KEYPOINTS
    
    build_fn = dataset_dict[task][dataset]
    testset = build_fn('test', config)
    testloader = torch.utils.data.DataLoader(testset, batch_size=1)

    pl_relpose = PL_RelPose.load_from_checkpoint(args.ckpt_path)
    pl_relpose.extractor = keypoint_dict[pl_relpose.hparams['features']](max_num_keypoints=test_num_keypoints, detection_threshold=0.0).eval().to(device)
    pl_relpose.module = pl_relpose.module.eval().to(device)

    preprocess_times, extract_times, regress_times = [], [], []
    adds, adis = [], []
    repr_errs = []
    R_errs, t_errs = [], []
    ts_errs = []
    results_dict = defaultdict(list)
    for i, data in enumerate(tqdm(testloader)):
        if dataset == 'ho3d' and args.obj_name is not None and data['objName'][0] != args.obj_name:
            continue
        image0, image1 = data['images'][0]
        K0, K1 = data['intrinsics'][0]
        T = torch.eye(4)
        T[:3, :3] = data['rotation'][0]
        T[:3, 3] = data['translation'][0]
        T = T.numpy()

        # with record_function("model_inference"):
        R_est, t_est, preprocess_time, extract_time, regress_time = pl_relpose.predict_one_data(data)
        preprocess_times.append(preprocess_time)
        extract_times.append(extract_time)
        regress_times.append(regress_time)

        t_err, R_err = relative_pose_error(T, R_est.cpu().numpy(), t_est.cpu().numpy(), ignore_gt_t_thr=0.0)

        R_errs.append(R_err)
        t_errs.append(t_err)

        ts_errs.append(torch.tensor(T[:3, 3] - t_est.cpu().numpy()).norm(2))

        if dataset == 'mapfree':
            repr_err = reprojection_error(R_est.cpu().numpy(), t_est.cpu().numpy(), T[:3, :3], T[:3, 3], K=K1, W=image1.shape[-1], H=image1.shape[-2])
            repr_errs.append(repr_err)
            R = R_est.detach().cpu().numpy()
            t = t_est.reshape(-1).detach().cpu().numpy()
            scene = data['scene_id'][0]
            estimated_pose = Pose(
                image_name=data['pair_names'][1][0],
                q=mat2quat(R).reshape(-1),
                t=t.reshape(-1),
                inliers=0
            )
            results_dict[scene].append(estimated_pose)

        if 'point_cloud' in data:
            adds.append(add(R_est.cpu().numpy(), t_est.cpu().numpy(), T[:3, :3], T[:3, 3], data['point_cloud'][0].numpy()))
            adis.append(adi(R_est.cpu().numpy(), t_est.cpu().numpy(), T[:3, :3], T[:3, 3], data['point_cloud'][0].numpy()))

    metrics = []
    values = []

    preprocess_times = np.array(preprocess_times) * 1000
    extract_times = np.array(extract_times) * 1000
    regress_times = np.array(regress_times) * 1000

    metrics.append('Extracting Time (ms)')
    values.append(f'{np.mean(extract_times):.1f}')
    
    metrics.append('Recovering Time (ms)')
    values.append(f'{np.mean(regress_times):.1f}')

    metrics.append('Total Time (ms)')
    values.append(f'{np.mean(extract_times) + np.mean(regress_times):.1f}')

    # ts_errs = np.array(ts_errs)
    # print(f'Median Trans. Error (m):\t{np.median(ts_errs):.2f}')
    # print(f'Median Rot. Error (°):\t{np.median(R_errs):.2f}')

    if task == 'object':
        metrics.append('Object ADD')
        values.append(f'{compute_continuous_auc(adds, np.linspace(0.0, 0.1, 1000)) * 100:.1f}')

        metrics.append('Object ADD-S')
        values.append(f'{compute_continuous_auc(adis, np.linspace(0.0, 0.1, 1000)) * 100:.1f}')

    if dataset == 'mapfree':
        re = np.array(repr_errs)

        metrics.append('VCRE @90px Prec.')
        values.append(f'{(re < 90).mean() * 100:.2f}')

        metrics.append('VCRE Med.')
        values.append(f'{np.median(re):.2f}')
        
        save_submission(results_dict, 'assets/new_submission.zip')


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('config', type=str, help='.yaml configure file path')
    parser.add_argument('ckpt_path', type=str)

    parser.add_argument('--device', type=str, default='cuda:0')
    parser.add_argument('--obj_name', type=str, default=None)

    return parser


if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    main(args)
